testing/python/language/test_tilelang_language_atomic_add.py (1)

375-376: Enable test: good; consider making it hardware-agnostic (float16) to avoid cc>=90 dependency

AtomicAddx4 for float32 is only selected when compute capability >= 90. On CI GPUs < 90 (e.g., A100 cc=80), this path may not vectorize and could cause flakiness for the tile-atomic path. Two options:

• Portable: call with float16 so vectorization is available broadly.
• Alternatively, gate/skip on device capability.

Apply this minimal change for portability:

-def test_tile_atomic_add():
-    run_tile_atomic_add(8, 128, 128, 32, 32)
+def test_tile_atomic_add():
+    run_tile_atomic_add(8, 128, 128, 32, 32, dtype="float16")

Also, consider removing or gating the debug prints in run_tile_atomic_add to keep test output clean (prints at Lines 58, 72, 73).

src/transform/atomicadd_vectorize.cc (4)

322-347: ParseIndex is too strict; accept const-expr strides and avoid false negatives

Requiring exactly one MulNode with a Var and an IntImm will miss common canonical forms:

• Stride may be a foldable const expr or come via casts (not a bare IntImm).
• Extra harmless multiplies like x*1 can appear pre-simplification.
• You only need a unique var*const match; other non-relevant muls shouldn’t invalidate the parse.

Refine by simplifying first, using as_const_int, and relaxing the check to “exactly one legal var*const mul” regardless of other muls:

-  auto ParseIndex = [](const PrimExpr &idx, PrimExpr &var_out,
-                       int &stride_out) -> bool {
+  auto ParseIndex = [](const PrimExpr &idx, PrimExpr &var_out,
+                       int &stride_out) -> bool {
     int mul_count = 0, legal_mul_count = 0;
     stride_out = -1;
     var_out = PrimExpr();
-    PostOrderVisit(idx, [&](const ObjectRef &obj) {
+    // Simplify to eliminate x*1 and fold-able constants.
+    arith::Analyzer az;
+    PrimExpr sidx = az.Simplify(idx);
+    PostOrderVisit(sidx, [&](const ObjectRef &obj) {
       if (const MulNode *mul = obj.as<MulNode>()) {
         mul_count++;
-        const VarNode *var = nullptr;
-        const IntImmNode *imm = nullptr;
-        if ((var = mul->a.as<VarNode>()) && (imm = mul->b.as<IntImmNode>())) {
-          var_out = mul->a;
-          stride_out = imm->value;
-          legal_mul_count++;
-        } else if ((var = mul->b.as<VarNode>()) &&
-                   (imm = mul->a.as<IntImmNode>())) {
-          var_out = mul->b;
-          stride_out = imm->value;
-          legal_mul_count++;
-        }
+        const VarNode *var = nullptr;
+        const int64_t *c = nullptr;
+        if ((var = mul->a.as<VarNode>()) && (c = as_const_int(mul->b))) {
+          var_out = mul->a;
+          stride_out = static_cast<int>(*c);
+          legal_mul_count++;
+        } else if ((var = mul->b.as<VarNode>()) && (c = as_const_int(mul->a))) {
+          var_out = mul->b;
+          stride_out = static_cast<int>(*c);
+          legal_mul_count++;
+        }
       }
     });
-    if (mul_count == 1 && legal_mul_count == 1)
-      return true;
-    return false;
+    return legal_mul_count == 1;
   };

Note: this uses as_const_int and simplification. If not already available, include tvm/arith/analyzer.h (already included).

---

362-368: Accumulate vectorize_size_max across multiple AtomicAdd sites

If the loop body contains multiple AtomicAdd calls, you currently overwrite vectorize_size_max. Prefer taking the max to avoid under-vectorizing later calls.

-          DataType dtype = bufload->dtype;
-          vectorize_size_max = GetVectorizeSizeMax(compute_capability, dtype);
+          DataType dtype = bufload->dtype;
+          vectorize_size_max = std::max(
+              vectorize_size_max, GetVectorizeSizeMax(compute_capability, dtype));

You’ll need:

• Add at top: #include

---

362-368: Guard against mis-identifying non-block vars as bx/by

ParseIndex will happily return any varconst (e.g., loop i1). Before accepting, assert the extracted vars are actual block indices (thread/block bindings) for safety, otherwise bail out. For example:

• Verify var_out.as() is bound in thread_binding as blockIdx.{x,y} (or matches expected bx/by symbols in this pass’ context).
• If that metadata isn’t available here, at least ensure both extracted vars differ and are not the loop var inside inner_for_.

This avoids rewriting with incorrect axes on more complex index expressions.

---

380-385: Extra sanity checks before rewriting

Before constructing the rewriter:

• Ensure bx_var and by_var are Vars: if (!bx_var.as() || !by_var.as()) return for_node;
• Optionally ensure stride_x > 0 && stride_y > 0.

This prevents emitting malformed truncdiv/truncmod expressions.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

testing/python/language/test_tilelang_language_atomic_add.py (1)

375-377: Remove noisy debug prints before re-enabling this test

Reactivating test_tile_atomic_add now exercises run_tile_atomic_add, which still contains print(kernel.get_kernel_source()) plus dumps of both 128×128 tensors. That’s hundreds of thousands of characters on every run and will swamp CI logs without adding assertion value. Please drop or gate those prints before merging.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

src/op/atomic_add.cc (1)

372-372: Consider removing or adjusting the log level.

The LOG(INFO) statement prints the vectorized loop IR to the console. This may be useful during development but could be noisy in production.

Consider one of the following:

1. Remove the log statement if it was added for debugging purposes only.
2. Change to VLOG(1) or a higher verbosity level to reduce noise in production logs.
3. If this is intentional diagnostic output, add a comment explaining why it's logged at INFO level.

Apply this diff to change to verbose logging:

-  LOG(INFO) << vectorized_thread_loop;
+  VLOG(1) << "Vectorized thread loop: " << vectorized_thread_loop;

src/transform/atomicadd_vectorize.cc (2)

33-55: Consider adding documentation for BufferIndiceSimplify.

The BufferIndiceSimplify class lacks documentation. Adding a brief comment explaining its purpose would improve maintainability.

Apply this diff to add documentation:

+/// \brief Simplifies buffer load and store indices using an analyzer.
+///
+/// This mutator visits BufferLoad and BufferStore nodes and simplifies
+/// their indices by applying the analyzer's Simplify method to each index.
 class BufferIndiceSimplify : public StmtExprMutator {

---

174-231: Consider adding documentation for the run() method.

The run() method implements complex loop transformation logic but lacks documentation explaining the transformation steps and the role of loop_layout and analyzer.

Apply this diff to add documentation:

+  /// \brief Transform and vectorize the for loop using the provided layout.
+  ///
+  /// \param for_node The original For loop to transform
+  /// \param loop_layout Fragment describing the loop layout transformation
+  /// \param analyzer Analyzer for simplifying indices and binding loop variables
+  /// \return Transformed and vectorized For loop
   For run(For for_node, const Fragment &loop_layout,
           arith::Analyzer *analyzer) {

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

src/op/atomic_add.cc (2)

375-383: Drop the shadowing GetArchInt lambda.

This lambda duplicates the file-scope GetArchInt(Target) function from src/target/utils.h (already included on Line 13). Using a local shadow risks divergence and bypasses any future validation fixes in the shared implementation.

Based on past review comments.

Apply this diff:

-  auto GetArchInt = [&](const Target &tgt) -> int {
-    int arch_int = 0;
-    if (auto s = tgt->GetAttr<String>("arch")) {
-      std::string arch = s.value();
-      if (arch.rfind("sm_", 0) == 0)
-        arch_int = std::stoi(arch.substr(3));
-    }
-    return arch_int;
-  };

And update line 455 and 506 to call GetArchInt(target) directly.

---

490-493: Propagate the dynamic predicate.

The planner's dynamic predicate is captured here but never used. When plan.dynamic is true and plan.condition is defined, the final vectorized_thread_loop should be wrapped in a guard (e.g., IfThenElse(pred, vectorized_thread_loop, thread_loop)) so the vectorized path executes only when the condition holds.

Based on past review comments.

If you intend to propagate the predicate, apply a diff similar to:

   auto vectorized_thread_loop =
       VectorizeAtomicAdd(thread_loop, GetArchInt(target));
+  if (ret.predicate.defined()) {
+    return IfThenElse(ret.predicate.value(), vectorized_thread_loop, thread_loop);
+  }
   return vectorized_thread_loop;

src/transform/atomicadd_vectorize.cc (1)

227-230: Pass the original PrimExprs to Downcast.

Downcast<BufferLoad>(node->args[1].as<BufferLoadNode>()) attempts to downcast a raw pointer, which fails to compile. Since Lines 220-224 already verify the args are BufferLoadNode*, pass the original PrimExpr objects directly to Downcast.

Based on past review comments.

Apply this diff:

-            const BufferLoad dst_node =
-                Downcast<BufferLoad>(node->args[1].as<BufferLoadNode>());
-            const BufferLoad value_node =
-                Downcast<BufferLoad>(node->args[2].as<BufferLoadNode>());
+            const BufferLoad dst_node = Downcast<BufferLoad>(node->args[1]);
+            const BufferLoad value_node = Downcast<BufferLoad>(node->args[2]);

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

src/transform/atomicadd_vectorize.h (1)

48-50: Remove unused planner state.

has_nonlocal_memory_access_ is introduced here but never referenced in atomicadd_vectorize.cc, so it will trip -Wunused-private-field (and become a hard error under -Werror). Please drop it or wire it into the implementation.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

src/op/atomic_add.cc (2)

372-380: Drop the shadowing GetArchInt lambda.

A file-scope GetArchInt(Target) already exists in src/target/utils.cc. Redefining an identical lambda here is redundant, risks divergence if the shared helper is updated with validation fixes, and bypasses the centralized implementation.

Replace the lambda with a call to the existing file-scope GetArchInt:

-  auto GetArchInt = [&](const Target &tgt) -> int {
-    int arch_int = 0;
-    if (auto s = tgt->GetAttr<String>("arch")) {
-      std::string arch = s.value();
-      if (arch.rfind("sm_", 0) == 0)
-        arch_int = std::stoi(arch.substr(3));
-    }
-    return arch_int;
-  };

And update the call sites (lines 452, 503) to use GetArchInt(target) directly.

---

487-504: Propagate planner predicate for dynamic vectorization.

AtomicAddVectorizePlanner::Plan can return dynamic=true with a condition guard. The predicate is captured at lines 487-490 but never used. When dynamic is true, the vectorized body must be wrapped with the guard before execution; otherwise the dynamic plan executes incorrectly.

Apply this diff to wrap the vectorized loop with the predicate when present:

   auto ret = AtomicAddInferLayout(transformed_loop,
                                   {T.target, T.thread_bounds, T.layout_map,
                                    analyzer, false, T.buffer_remap});
   Fragment loop_layout = ret.loop_layout;
   auto thread_loop =
       PartitionLoop(transformed_loop, T.thread_var, analyzer, loop_layout);
   auto vectorized_thread_loop =
       VectorizeAtomicAdd(thread_loop, GetArchInt(target));
-  return vectorized_thread_loop;
+  if (ret.predicate.defined()) {
+    return IfThenElse(ret.predicate.value(), vectorized_thread_loop);
+  }
+  return vectorized_thread_loop;
 }

src/op/atomic_add.cc (1)

398-409: Consider adding consistency checks for buffer indices.

Unlike ParallelLoopNestVisitor (see src/op/parallel.cc lines 147-172), AtomicLoopNestCollector does not verify that the same buffer is accessed with structurally equal indices across multiple access sites. If a buffer has inconsistent access patterns, the collector silently overwrites the indices, potentially masking layout inference bugs.

Apply this diff to add consistency validation:

 void VisitStmt_(const BufferStoreNode *op) final {
   if (op->buffer.scope() == "local.fragment") {
+    if (indice_map.count(op->buffer)) {
+      ICHECK(StructuralEqual()(indice_map[op->buffer], op->indices))
+          << op->buffer << ": " << op->indices << " vs "
+          << indice_map[op->buffer];
+    }
     indice_map.Set(op->buffer, op->indices);
     writes.insert(op->buffer);
   }
   StmtExprVisitor::VisitStmt_(op);
 }
 void VisitExpr_(const BufferLoadNode *op) final {
   if (op->buffer.scope() == "local.fragment") {
+    if (indice_map.count(op->buffer)) {
+      ICHECK(StructuralEqual()(indice_map[op->buffer], op->indices))
+          << op->buffer << ": " << op->indices << " vs "
+          << indice_map[op->buffer];
+    }
     indice_map.Set(op->buffer, op->indices);
   }
   StmtExprVisitor::VisitExpr_(op);
 }

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

src/op/atomic_add.cc (2)

378-386: Drop the redundant local GetArchInt lambda.

This local lambda duplicates the file-scope GetArchInt(Target) function already available in ../target/utils.h (included at line 13). The duplication creates maintenance risk and bypasses any future improvements to the shared implementation.

Apply this diff to remove the lambda and use the file-scope function:

-  auto GetArchInt = [&](const Target &tgt) -> int {
-    int arch_int = 0;
-    if (auto s = tgt->GetAttr<String>("arch")) {
-      std::string arch = s.value();
-      if (arch.rfind("sm_", 0) == 0)
-        arch_int = std::stoi(arch.substr(3));
-    }
-    return arch_int;
-  };

Then update the call sites at lines 458 and 509 to use the file-scope GetArchInt(target) directly (note: the function is already in scope via the include at line 13).

---

502-510: Propagate the planner predicate for dynamic vectorization.

AtomicAddInferLayout captures a predicate from the planner when dynamic=true (lines 493-496) but this predicate is never used. When dynamic vectorization is planned, the vectorized body should be wrapped with the predicate guard to ensure correct execution.

Apply this diff to wrap the vectorized loop with the predicate when present:

   auto ret = AtomicAddInferLayout(transformed_loop,
                                   {T.target, T.thread_bounds, T.layout_map,
                                    analyzer, false, T.buffer_remap});
   Fragment loop_layout = ret.loop_layout;
   auto thread_loop =
       PartitionLoop(transformed_loop, T.thread_var, analyzer, loop_layout);
   auto vectorized_thread_loop =
       VectorizeAtomicAdd(thread_loop, GetArchInt(target));
-  return vectorized_thread_loop;
+  if (ret.predicate.defined()) {
+    return IfThenElse(ret.predicate.value(), vectorized_thread_loop);
+  } else {
+    return vectorized_thread_loop;
+  }

src/op/atomic_add.cc (1)

388-417: Add index consistency validation in AtomicLoopNestCollector.

The collector sets indice_map without validating that repeated accesses to the same buffer use structurally equal indices. This differs from ParallelLoopNestVisitor (in src/op/parallel.cc lines 147-172), which validates consistency with ICHECK(StructuralEqual()(...). Silently overwriting indices can hide bugs where the same buffer is accessed with different patterns.

Apply this diff to add validation:

     void VisitStmt_(const BufferStoreNode *op) final {
       if (op->buffer.scope() == "local.fragment") {
+        if (indice_map.count(op->buffer)) {
+          ICHECK(StructuralEqual()(indice_map[op->buffer], op->indices))
+              << op->buffer << ": " << op->indices << " and "
+              << indice_map[op->buffer];
+        }
         indice_map.Set(op->buffer, op->indices);
         writes.insert(op->buffer);
       }
       StmtExprVisitor::VisitStmt_(op);
     }
     void VisitExpr_(const BufferLoadNode *op) final {
       if (op->buffer.scope() == "local.fragment") {
+        if (indice_map.count(op->buffer)) {
+          ICHECK(StructuralEqual()(indice_map[op->buffer], op->indices))
+              << op->buffer << ": " << op->indices << " and "
+              << indice_map[op->buffer];
+        }
         indice_map.Set(op->buffer, op->indices);
       }
       StmtExprVisitor::VisitExpr_(op);
     }

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

src/op/atomic_add.cc (2)

412-420: Remove the redundant GetArchInt lambda.

A file-scope GetArchInt(Target) function already exists (imported via #include "../target/utils.h" at line 13, defined in /src/target/utils.cc). This local lambda duplicates that logic and shadows the shared implementation. Use the existing function directly instead.

Apply this diff:

-  auto GetArchInt = [&](const Target &tgt) -> int {
-    int arch_int = 0;
-    if (auto s = tgt->GetAttr<String>("arch")) {
-      std::string arch = s.value();
-      if (arch.rfind("sm_", 0) == 0)
-        arch_int = std::stoi(arch.substr(3));
-    }
-    return arch_int;
-  };

Then update line 492 and 543 to call the file-scope function:

-    int sm = GetArchInt(target);
+    int sm = tl::GetArchInt(target);

-      VectorizeAtomicAdd(thread_loop, GetArchInt(target));
+      VectorizeAtomicAdd(thread_loop, tl::GetArchInt(target));

---

536-544: Missing predicate guard for dynamic vectorization.

The planner's predicate (captured at line 533) is never used. When plan.dynamic is true, the vectorized body must be wrapped with the predicate guard to ensure correctness. Currently, line 544 returns the vectorized loop unconditionally, which can produce incorrect results for dynamic plans.

Apply this diff to wrap the vectorized loop when a predicate exists:

   auto ret = AtomicAddInferLayout(transformed_loop,
                                   {T.target, T.thread_bounds, T.layout_map,
                                    analyzer, false, T.buffer_remap});
   Fragment loop_layout = ret.loop_layout;
   auto thread_loop =
       PartitionLoop(transformed_loop, T.thread_var, analyzer, loop_layout);
   auto vectorized_thread_loop =
       VectorizeAtomicAdd(thread_loop, tl::GetArchInt(target));
-  return vectorized_thread_loop;
+  if (ret.predicate.defined()) {
+    return IfThenElse(ret.predicate.value(), vectorized_thread_loop);
+  }
+  return vectorized_thread_loop;
 }

src/op/atomic_add.cc (2)

477-490: Consider renaming read_src for clarity.

The variable read_src (line 477) doesn't verify whether the buffer is read-only; it simply selects the buffer with the highest rank from indice_map. Consider renaming it to layout_src or infer_src to better reflect its purpose.

---

523-524: Remove unnecessary remapped variable.

Line 523 assigns loop to remapped without any transformation. This appears to be leftover code from a refactoring. Simplify by using loop directly at line 524.

Apply this diff:

-      const For &remapped = loop;
-      loop_layout = PlanLoopPartition(remapped, vec, args.thread_bounds);
+      loop_layout = PlanLoopPartition(loop, vec, args.thread_bounds);

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro